import os
import torch
import torch.utils.data as data
import torch
import torchvision.transforms as transforms
import random
from albumentations import Compose, OneOf
from PIL import Image, ImageOps
import numpy as np
from . import preprocess
from .stereo_albumentation import RandomShiftRotate, GaussNoiseStereo, RGBShiftStereo, \
    RandomBrightnessContrastStereo, random_crop, horizontal_flip
from . import transforms
from .transforms import RandomColor
import cv2
import pdb

from ruamel.yaml import YAML


configs = [
    'cfg_coex.yaml',
    'cfg_psm.yaml'
    ]
config_num = 0


IMG_EXTENSIONS = [
    '.jpg', '.JPG', '.jpeg', '.JPEG',
    '.png', '.PNG', '.ppm', '.PPM', '.bmp', '.BMP',
]

def is_image_file(filename):
    return any(filename.endswith(extension) for extension in IMG_EXTENSIONS)

def load_configs(path):
    cfg = YAML().load(open(path, 'r'))
    backbone_cfg = YAML().load(
        open(cfg['model']['stereo']['backbone']['cfg_path'], 'r'))
    cfg['model']['stereo']['backbone'].update(backbone_cfg)
    return cfg

def dataloader(filepath, trainAll=False):
    cfg = load_configs('./configs/stereo/{}'.format(configs[config_num]))
    kfold = cfg['training']['kfold']
    if kfold != 0:
        train_list_path = 'dataloaders/stereo/lists/usvinland_train_'+str(kfold)+'.list'
        test_list_path = 'dataloaders/stereo/lists/usvinland_val_'+str(kfold)+'.list'
    else:
        train_list_path = 'dataloaders/stereo/lists/usvinland_train.list'
        test_list_path = 'dataloaders/stereo/lists/usvinland_val.list'
    train_list_f = open(train_list_path, 'r')
    test_list_f = open(test_list_path, 'r')
    print('kfold:', kfold)
    print(train_list_path)
    print(test_list_path)
    train_list_ = train_list_f.readlines()
    test_list_ = test_list_f.readlines()
    if trainAll:
        train_list = train_list_ + test_list_
        test_list = train_list_ + test_list_
    else:
        train_list = train_list_
        test_list = test_list_

    left_train = []
    right_train = []
    disp_train_L = []
    for i in range(len(train_list)):
        name = train_list[i].split('.')[0] + '.jpg'
        left_train.append(filepath + '/Left_Img_Rectified/' + name)
        right_train.append(filepath + '/Right_Img_Rectified/' + name)
        disp_train_L.append(filepath + '/Disp_Map/' + name.replace('.jpg', '.png'))

    left_val = []
    right_val = []
    disp_val_L = []
    for i in range(len(test_list)):
        name = test_list[i].split('.')[0] + '.jpg'
        left_val.append(filepath + '/Left_Img_Rectified/' + name)
        right_val.append(filepath + '/Right_Img_Rectified/' + name)
        disp_val_L.append(filepath + '/Disp_Map/' + name.replace('.jpg', '.png'))

    return left_train, right_train, disp_train_L, left_val, right_val, disp_val_L

def default_loader(path):
    return cv2.imread(path)
    #return Image.open(path).convert('RGB')

def disparity_loader(path):
    return Image.open(path)


class ImageLoader(data.Dataset):
    def __init__(self, left, right, left_disparity, th=320, tw=640, shift=0, training=True, loader=default_loader, dploader= disparity_loader):

        self.left = left
        self.right = right
        self.disp_L = left_disparity
        self.loader = loader
        self.dploader = dploader
        self.training = training
        self.th = th
        self.tw = tw
        self.shift = shift

    def __getitem__(self, index):
        batch = dict()

        left  = self.left[index]
        right = self.right[index]
        disp_L= self.disp_L[index]

        left_img = self.loader(left)
        right_img = self.loader(right)
        dataL = self.dploader(disp_L)

        dataL = np.ascontiguousarray(dataL,dtype=np.float32)/255*50

        if self.training:
            pad_h, pad_w = 320-left_img.shape[0], 640-left_img.shape[1]

            left_img = np.pad(left_img,((0,pad_h),(0,pad_w),(0,0)))
            right_img = np.pad(right_img,((0,pad_h),(0,pad_w),(0,0)))

            h,w,_ = left_img.shape
            th, tw = self.th, self.tw

            shift = random.randint(-self.shift,self.shift)
            x1 = random.randint(0, w - tw)
            y1 = random.randint(0, h - th)

            # if x1 + shift < 0 or  x1 + shift + tw > w:
            shift = 0

            left_img_raw = left_img[y1:y1+th,x1+shift:x1+shift+tw,:]
            right_img_raw = right_img[y1:y1+th,x1:x1+tw,:]

            imL_lab = cv2.cvtColor(
                left_img_raw,#cv2.resize(left_img,None,None,0.25,0.25),
                cv2.COLOR_BGR2LAB)

            dataL = np.pad(dataL[:,:,np.newaxis],((0,pad_h),(0,pad_w),(0,0)))[:,:,0]
            dataL = dataL[y1:y1 + th, x1 + shift:x1 + tw + shift]
            dataL = dataL - shift

            img = {'left':left_img_raw,'right':right_img_raw}
            # img = self.train_aug(img)

            left_img_raw, right_img_raw = img['left'], img['right']

            processed = preprocess.get_transform(augment=False)
            left_img   = processed(left_img_raw)
            right_img  = processed(right_img_raw)

            left_img_raw = np.transpose(left_img_raw,(2,0,1)).astype(np.float32)
            right_img_raw = np.transpose(right_img_raw,(2,0,1)).astype(np.float32)

            batch['imgL'], batch['imgR'], batch['disp_true'] = left_img, right_img, dataL
            batch['imgLRaw'], batch['imgRRaw'], batch['imgLLab'] = left_img_raw, right_img_raw, imL_lab
            batch['x1'], batch['y1'] = x1, y1

            return batch
        else:
            h,w,_ = left_img.shape
            imL = left_img
            pad_h, pad_w = 320-h, 640-w

            left_img_raw = left_img  # np.pad(left_img,((0,pad_h),(0,pad_w),(0,0)))
            right_img_raw = right_img  # np.pad(right_img,((0,pad_h),(0,pad_w),(0,0)))

            imL_lab = cv2.cvtColor(
                left_img_raw,#cv2.resize(left_img,None,None,0.25,0.25),
                cv2.COLOR_BGR2LAB)

            processed = preprocess.get_transform(augment=False)
            left_img       = processed(left_img_raw)
            right_img      = processed(right_img_raw)

            batch['imgL'], batch['imgR'], batch['disp_true'] = left_img, right_img, dataL
            batch['imgLLab'], batch['dataname'], batch['filename'] = imL_lab, 'usvinland', left.split('/')[-1].replace('.jpg', '')

            return batch

    def __len__(self):
        return len(self.left)

    def train_aug(self, img):
        transformation = Compose([
                RGBShiftStereo(always_apply=True, p_asym=0.5),
                RandomBrightnessContrastStereo(always_apply=True, p_asym=0.5)
                ])
        return transformation(**img)

        # transformation = transforms.Compose([
        #         RandomColor()
        #         ])
        # return transformation(img)
